import torch
import torchvision.models as models
import torch.nn as nn
import torchvision.transforms as transforms
from AI_Densenet.my_dataloader import My_ImageFolder,Compoment_DataLoader
from torchvision import  utils as vutils

class Compoment_Classify():
    def __init__(self):
        super(Compoment_Classify,self).__init__()
        self.classes = ['C', 'CN', 'D', 'EG', 'IC', 'L', 'Other', 'Other_NG', 'P', 'Q', 'R', 'X', '空焊盘']
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.num_class = len(self.classes)
        self.batch_sise = 8
        self.densenet = models.densenet161()
        self.weight = r'E:\DenseNet\result_3d_Pemtron\3d_Pemtron_densenet_161_2.pth'
        num_ftrs = self.densenet.classifier.in_features
        self.densenet.classifier = nn.Linear(num_ftrs,self.num_class)
        self.densenet.load_state_dict(torch.load(self.weight,map_location=self.device))
        self.transform = transforms.Compose(
            [transforms.ToTensor()]
        )
    def start_det(self,image_rows):
        self.densenet.to(self.device)
        self.densenet.eval()
        dataset = Compoment_DataLoader(image_rows,transform=self.transform)
        dataloader = torch.utils.data.DataLoader(dataset,batch_size = self.batch_sise,shuffle=False,num_workers=0)
        outputs = {}
        with torch.no_grad():
            for i,(data) in enumerate(dataloader):
                keys = data[0]
                images = data[1]
                images = images.to(self.device)
                output = self.densenet(images)
                output = torch.softmax(output, dim=1)

                _, pred = output.topk(1, 1, True, True)
                for j,key in enumerate(keys):
                    outputs[key] = self.classes[int(pred[j][0])]
                # print(pred)
                # outputs.extend(pred)
        # print(outputs)
        return outputs
        pass